#!/usr/bin/env python3
"""
代码优化器MCP服务器

这个脚本实现了一个MCP服务器，允许通过MCP协议使用代码优化器功能。
"""

import os
import sys
import json
from pathlib import Path
from typing import Dict, Any, List, Optional

# 确保代码优化器包可以被导入
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

try:
    from code_optimizer.core import CodeOptimizer, DEFAULT_CONFIG
    from modelcontextprotocol.sdk.server import Server
    from modelcontextprotocol.sdk.server.stdio import StdioServerTransport
    from modelcontextprotocol.sdk.types import (
        CallToolRequestSchema,
        ErrorCode,
        ListToolsRequestSchema,
        McpError,
    )
except ImportError as e:
    print(f"错误：无法导入必要的包: {e}")
    print("请确保已安装所需的依赖:")
    print("  pip install @modelcontextprotocol/sdk")
    sys.exit(1)

class CodeOptimizerServer:
    """代码优化器MCP服务器类"""
    
    def __init__(self):
        self.server = Server(
            {
                "name": "code-optimizer-server",
                "version": "0.1.0",
            },
            {
                "capabilities": {
                    "resources": {},
                    "tools": {},
                }
            }
        )
        
        self.optimizer = CodeOptimizer(DEFAULT_CONFIG.copy())
        
        self.setup_tool_handlers()
        
        # 错误处理
        self.server.onerror = lambda error: print(f"[MCP错误] {error}", file=sys.stderr)
        
        # 处理中断信号
        import signal
        signal.signal(signal.SIGINT, self.handle_interrupt)
    
    def handle_interrupt(self, sig, frame):
        """处理中断信号"""
        print("\n正在关闭MCP服务器...", file=sys.stderr)
        self.server.close()
        sys.exit(0)
    
    def setup_tool_handlers(self):
        """设置工具处理器"""
        self.server.setRequestHandler(ListToolsRequestSchema, self.handle_list_tools)
        self.server.setRequestHandler(CallToolRequestSchema, self.handle_call_tool)
    
    async def handle_list_tools(self, request):
        """处理工具列表请求"""
        return {
            "tools": [
                {
                    "name": "optimize_code",
                    "description": "优化代码文件或目录",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "path": {
                                "type": "string",
                                "description": "要优化的文件或目录路径"
                            },
                            "split_threshold": {
                                "type": "number",
                                "description": "文件拆分阈值（行数）",
                                "default": DEFAULT_CONFIG["split_threshold"]
                            },
                            "add_async": {
                                "type": "boolean",
                                "description": "是否添加异步支持",
                                "default": DEFAULT_CONFIG["add_async"]
                            },
                            "optimize": {
                                "type": "boolean",
                                "description": "是否优化代码",
                                "default": DEFAULT_CONFIG["optimize"]
                            },
                            "monitor": {
                                "type": "boolean",
                                "description": "是否添加执行时间监控",
                                "default": DEFAULT_CONFIG["monitor"]
                            }
                        },
                        "required": ["path"]
                    }
                },
                {
                    "name": "analyze_code",
                    "description": "分析代码文件或目录",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "path": {
                                "type": "string",
                                "description": "要分析的文件或目录路径"
                            }
                        },
                        "required": ["path"]
                    }
                }
            ]
        }
    
    async def handle_call_tool(self, request):
        """处理工具调用请求"""
        tool_name = request.params.name
        args = request.params.arguments
        
        if tool_name == "optimize_code":
            return await self.handle_optimize_code(args)
        elif tool_name == "analyze_code":
            return await self.handle_analyze_code(args)
        else:
            raise McpError(
                ErrorCode.MethodNotFound,
                f"未知工具: {tool_name}"
            )
    
    async def handle_optimize_code(self, args):
        """处理代码优化请求"""
        if not isinstance(args, dict) or "path" not in args:
            raise McpError(
                ErrorCode.InvalidParams,
                "缺少必要参数: path"
            )
        
        path = args["path"]
        
        # 更新配置
        config = DEFAULT_CONFIG.copy()
        for key in ["split_threshold", "add_async", "optimize", "monitor"]:
            if key in args:
                config[key] = args[key]
        
        # 创建优化器
        self.optimizer = CodeOptimizer(config)
        
        # 处理路径
        path_obj = Path(path)
        results = []
        
        try:
            if path_obj.is_file():
                success = self.optimizer.optimize_file(str(path_obj))
                results.append({
                    "file": str(path_obj),
                    "success": success
                })
            elif path_obj.is_dir():
                for root, _, files in os.walk(path_obj):
                    for file in files:
                        file_path = Path(root) / file
                        if not self.optimizer.should_ignore(str(file_path)):
                            success = self.optimizer.optimize_file(str(file_path))
                            results.append({
                                "file": str(file_path),
                                "success": success
                            })
            else:
                return {
                    "content": [
                        {
                            "type": "text",
                            "text": f"路径不存在: {path}"
                        }
                    ],
                    "isError": True
                }
            
            # 统计结果
            total = len(results)
            successful = sum(1 for r in results if r["success"])
            
            return {
                "content": [
                    {
                        "type": "text",
                        "text": f"优化完成: 共处理 {total} 个文件，成功 {successful} 个，失败 {total - successful} 个"
                    },
                    {
                        "type": "json",
                        "json": results
                    }
                ]
            }
            
        except Exception as e:
            return {
                "content": [
                    {
                        "type": "text",
                        "text": f"优化过程中出错: {str(e)}"
                    }
                ],
                "isError": True
            }
    
    async def handle_analyze_code(self, args):
        """处理代码分析请求"""
        if not isinstance(args, dict) or "path" not in args:
            raise McpError(
                ErrorCode.InvalidParams,
                "缺少必要参数: path"
            )
        
        path = args["path"]
        path_obj = Path(path)
        
        try:
            if not path_obj.exists():
                return {
                    "content": [
                        {
                            "type": "text",
                            "text": f"路径不存在: {path}"
                        }
                    ],
                    "isError": True
                }
            
            analysis_results = self._analyze_path(path_obj)
            
            return {
                "content": [
                    {
                        "type": "text",
                        "text": f"代码分析完成: {path}"
                    },
                    {
                        "type": "json",
                        "json": analysis_results
                    }
                ]
            }
            
        except Exception as e:
            return {
                "content": [
                    {
                        "type": "text",
                        "text": f"分析过程中出错: {str(e)}"
                    }
                ],
                "isError": True
            }
    
    def _analyze_path(self, path: Path) -> Dict[str, Any]:
        """分析文件或目录"""
        if path.is_file():
            return self._analyze_file(path)
        
        results = {
            "path": str(path),
            "type": "directory",
            "files": []
        }
        
        for root, _, files in os.walk(path):
            for file in files:
                file_path = Path(root) / file
                if not self.optimizer.should_ignore(str(file_path)):
                    file_analysis = self._analyze_file(file_path)
                    results["files"].append(file_analysis)
        
        return results
    
    def _analyze_file(self, path: Path) -> Dict[str, Any]:
        """分析单个文件"""
        suffix = path.suffix
        
        if suffix not in self.optimizer.config["language_handlers"]:
            return {
                "path": str(path),
                "type": "unsupported",
                "size": path.stat().st_size,
                "lines": None
            }
        
        try:
            with open(path, 'r', encoding='utf-8') as f:
                content = f.read()
                lines = content.splitlines()
            
            language = self.optimizer.config["language_handlers"][suffix]
            handler = self.optimizer._get_handler(path)
            handler.load()
            handler.analyze()
            
            return {
                "path": str(path),
                "type": "file",
                "language": language,
                "size": path.stat().st_size,
                "lines": len(lines),
                "needs_splitting": len(lines) > self.optimizer.config["split_threshold"],
                "imports": list(handler.imports) if hasattr(handler, "imports") else [],
                "functions": handler.functions if hasattr(handler, "functions") else [],
                "classes": handler.classes if hasattr(handler, "classes") else []
            }
            
        except Exception as e:
            return {
                "path": str(path),
                "type": "error",
                "error": str(e)
            }
    
    async def run(self):
        """运行MCP服务器"""
        transport = StdioServerTransport()
        await self.server.connect(transport)
        print("代码优化器MCP服务器已启动", file=sys.stderr)

def main():
    """主函数"""
    server = CodeOptimizerServer()
    server.run()

if __name__ == "__main__":
    main()